#pragma once
#include "NvInferRuntime.h"
#include "cuda_fp16.h"

namespace nvinfer1 {

template <DataType dtype> struct DataTypeTraits {};

template <> struct DataTypeTraits<DataType::kFLOAT> {
  static constexpr size_t size = sizeof(float);
};

template <> struct DataTypeTraits<DataType::kHALF> {
  static constexpr size_t size = sizeof(half);
};

template <> struct DataTypeTraits<DataType::kINT8> {
  static constexpr size_t size = sizeof(int8_t);
};

template <> struct DataTypeTraits<DataType::kINT32> {
  static constexpr size_t size = sizeof(int32_t);
};

template <> struct DataTypeTraits<DataType::kBOOL> {
  static constexpr size_t size = sizeof(bool);
};

} // namespace nvinfer1
